package cc.shacocloud.luanniao.core.service.ai.doubao

import com.volcengine.ark.runtime.model.completion.chat.*
import com.volcengine.ark.runtime.model.embeddings.EmbeddingResult
import org.springframework.ai.chat.messages.AssistantMessage
import org.springframework.ai.chat.messages.MessageType
import org.springframework.ai.chat.messages.MessageType.*
import org.springframework.ai.chat.metadata.ChatGenerationMetadata
import org.springframework.ai.chat.metadata.ChatResponseMetadata
import org.springframework.ai.chat.metadata.EmptyUsage
import org.springframework.ai.chat.model.ChatResponse
import org.springframework.ai.chat.model.Generation
import org.springframework.ai.chat.prompt.Prompt
import org.springframework.ai.embedding.Embedding
import org.springframework.ai.embedding.EmbeddingResponse
import org.springframework.ai.embedding.EmbeddingResponseMetadata

fun ChatCompletionChunk.toChatResponse(): ChatResponse {
    val meta = ChatResponseMetadata.builder()
        .withId(id)
        .withModel(model)
        .withUsage(usage?.let { DoubaoUsage(it) } ?: EmptyUsage())
        .build()

    val generations = choices.map { it.toGeneration() }
    return ChatResponse(generations, meta)
}

fun ChatCompletionResult.toChatResponse(): ChatResponse {
    val meta = ChatResponseMetadata.builder()
        .withId(id)
        .withModel(model)
        .withUsage(DoubaoUsage(usage))
        .build()

    val generations = choices.map { it.toGeneration() }
    return ChatResponse(generations, meta)
}

fun Prompt.toMessages(): List<ChatMessage> {
    return instructions.map {
        ChatMessage.builder()
            .role(it.messageType.toChatMessageRole())
            .content(it.content)
            .build()
    }
}

fun ChatCompletionChoice.toGeneration(): Generation {
    return Generation(
        AssistantMessage(
            message.stringContent(),
            emptyMap(),
            message.toolCalls?.map { it.toToolCall() } ?: emptyList()),
        ChatGenerationMetadata.from(finishReason, null)
    )
}

fun ChatToolCall.toToolCall(): AssistantMessage.ToolCall {
    return AssistantMessage.ToolCall(
        id, type, function.name, function.arguments
    )
}

fun MessageType.toChatMessageRole(): ChatMessageRole {
    return when (this) {
        USER -> ChatMessageRole.USER
        ASSISTANT -> ChatMessageRole.ASSISTANT
        SYSTEM -> ChatMessageRole.SYSTEM
        TOOL -> ChatMessageRole.TOOL
    }
}

fun EmbeddingResult.toEmbeddingResponse(): EmbeddingResponse {
    val embeddings = data.map { it.toEmbedding() }
    val metadata = EmbeddingResponseMetadata(model, usage?.let { DoubaoUsage(it) } ?: EmptyUsage())
    return EmbeddingResponse(embeddings, metadata)
}

fun com.volcengine.ark.runtime.model.embeddings.Embedding.toEmbedding(): Embedding {
    val embedding = FloatArray(embedding.size) {
        embedding[it].toFloat()
    }
    return Embedding(embedding, index)
}
